import os
import socket
import threading
import queue
import time

DEBUG_ENABLE = True
def debug(str):
    if DEBUG_ENABLE: print("UDP Server>>" + str)

def current_time_ms():
    return time.time_ns()/1000000

def get_name(addr):
    return '%s-%d' % (addr[0], addr[1])

### udp 服务器监听任务，将客服端保存到服务器客户端表
def udp_server_rec_task(ts):
    while True:
        try:
            if ts is not None:
                data, addr = ts.server_socket.recvfrom(4096)
                cl = udp_server_client(ts, addr)
                cl = ts.add_client_map_ele(addr, cl)
                ts.put_rec_queue((cl.name, data))
            else:
                debug("Error: udp Server Is None!")
                return
        except:
            debug("Error: udp Server Fail!")
            #return

### udp 服务器分发任务，将所有客户端接收到的数据分发到各个客户端
def udp_server_send_task(ts):
    while True:
        while ts.has_rec_data():
            ele = ts.get_a_rec_data()
            name = ele[0]
            data = ele[1]
            ts.client_map_lock.acquire()
            keys = list(ts.client_map.keys())
            for k in keys:
                if ts.auto_test_enable(): #自动回测
                    if k==name:
                        ts.client_map[k].put_send_queue(data)
                else:
                    if k!=name:
                        ts.client_map[k].put_send_queue(data) #往客户端缓冲区中写入待发送的数据
            ts.client_map_lock.release()
        ts.update_client_map()
        time.sleep(0.0001)

### udp 客户端发送数据任务
def udp_server_client_send_task(ts, cl):
    while True:
        try:
            cl.send_data()
        except:
            ts.delete_client_map_ele(cl.name)
            debug("Delete Client: %s (%d)" % (cl.name, ts.get_client_count()))
            return
        time.sleep(0.0001)

### udp 客户端类
class udp_server_client(object):
    def __init__(self, server, addr):
        self.server = server
        self.name = get_name(addr)
        self.client_addr = addr
        self.rec_queue = queue.Queue(0)
        self.send_queue = queue.Queue(0)
        self.last_rec_time_ms = current_time_ms()

        threading.Thread(target=udp_server_client_send_task, args=(self.server,self)).start()

    def __has_send_data(self):
        return (not self.send_queue.empty())

    def send_data(self):
        while self.__has_send_data():
            self.server.server_socket.sendto(self.send_queue.get(), self.client_addr)

    def put_send_queue(self, data):
        self.send_queue.put(data)

    def put_rec_queue(self, data):
        self.rec_queue.put(data)

    def recv(self, len):
        try:
            data = self.client_socket.recv(len)
            self.last_rec_time_ms = current_time_ms()
            return data
        except:
            pass

### udp服务器类
class udp_server(object):
    def __init__(self, addr='127.0.0.1', port=9999, client_timeout=10, auto_test_en=False):
        self.server_socket = None
        self.addr = addr
        self.port = port
        self.auto_test_en = auto_test_en
        self.client_time_out = client_timeout*1000
        self.client_map_lock = threading.Lock()
        # udp 服务器客户端列表
        self.client_map = {} #name: client
        # udp 将服务器所有客户端接受到的数据保存到队列
        self.rec_queue = queue.Queue(0) #(name, data)

    def add_client_map_ele(self, addr, cl):
        name = get_name(addr)
        now = current_time_ms()
        self.client_map_lock.acquire()
        keys = list(self.client_map.keys())
        if len(keys)<=0 or (name not in keys):
            self.client_map[name] = cl
            debug('%d: add client %s' % (len(list(self.client_map.keys())), name))
        else:
            self.client_map[name].last_rec_time_ms = now
        self.client_map_lock.release()
        return self.client_map[name]

    def delete_client_map_ele(self, name):
        self.client_map_lock.acquire()
        if name in self.client_map.keys():
            del self.client_map[name]
        self.client_map_lock.release()

    def has_rec_data(self):
        return (not self.rec_queue.empty())

    def get_a_rec_data(self):
        return self.rec_queue.get()

    def put_rec_queue(self, ele):
        if (ele is not None) and (ele[1] is not None) and len(ele[1])>0:
            self.rec_queue.put(ele)

    def auto_test_enable(self):
        return self.auto_test_en

    def update_client_map(self):
        now = current_time_ms()
        self.client_map_lock.acquire()
        keys = list(self.client_map.keys())
        for k in keys:
            lt = self.client_map[k].last_rec_time_ms
            if now>lt and now-lt>=self.client_time_out:
                del self.client_map[k]
                debug("Timeout: Delete Client %s" % k)
        self.client_map_lock.release()
        if len(list(self.client_map.keys()))<=0:
            if self.has_rec_data():
                self.rec_queue.queue.clear()
                debug("Clear Server Rec Queue")

    def start(self):
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.server_socket.bind((self.addr, self.port))

        threading.Thread(target=udp_server_rec_task, args=(self,)).start()
        threading.Thread(target=udp_server_send_task, args=(self,)).start()

if __name__=='__main__':
    try:
        udp_server('192.168.10.243', 9999,  auto_test_en=False).start() #分发服务器
        udp_server('192.168.10.243', 14550, auto_test_en=True).start() #自动回测服务器
        while True:
            time.sleep(0.1)
    except:

        os._exit(0)

